#include "main.h"

IPandDomain DNS_table[AMOUNT];//DNS域名解析表
IDchange IDTransTable[AMOUNT];//ID转换表
int IDCount;//ID转换表中的条目个数
char URL[LENGTH];//域名

int main(int argc, char *argv[]) {
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);
    SOCKET sockServer, sockLocal;//外部DNS套接字, 本地DNS套接字
    struct sockaddr_in serverAddr, localAddr, clientAddr;//外部DNS地址, 本地DNS地址, 客户端地址
    char recvBuf[BUF_SIZE], sendBuf[BUF_SIZE];//接收缓存, 发送缓存
    char tablePath[512];//解析表路径
    char server_DNS[64];//域名服务器IP
    int entries_num;//域名解析表条目数
    int level;

    //调试级别0
    if (argc == 1) {
        printf("debug level is 0\n");
        strcpy(tablePath, PATH);//获取自己读取文件的地址
        strcpy(server_DNS, SERVER_DNS_IP);//获取名字服务器IP
        level = 0;
    }
        //调试级别1
    else if (argc == 4 && strcmp(argv[1], "-d") == 0) {
        printf("debug level is 1\n");
        strcpy(tablePath, argv[3]);
        strcpy(server_DNS, argv[2]);
        level = 1;
    }
        //调试级别2
    else if (argc == 3 && strcmp(argv[1], "-dd") == 0) {
        printf("debug level is 2\n");
        strcpy(tablePath, PATH);
        strcpy(server_DNS, argv[2]);
        level = 2;
    }
        //参数不匹配
    else {
        printf("Argument is incorrect!\n");
        printf("Command line format is: dnsrelay [-d | -dd] [dns-server-ipAddr] [filename]\n");
        printf("The file path cannot be spaced!\n");
        exit(1);
    }

    entries_num = getTable(tablePath);//获取解析表
    //初始化ID转换表
    for (int i = 0; i < AMOUNT; i++) {
        IDTransTable[i].oldID = 0;
        memset(&(IDTransTable[i].client), 0, sizeof(SOCKADDR_IN));
    }

    // 初始套接字
    printf("Initializing socket...\n");
    // 创建本地DNS套接字和外部DNS套接字
    sockLocal = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    memset(&localAddr, 0, sizeof(localAddr));
    localAddr.sin_family = AF_INET;
    localAddr.sin_addr.S_un.S_addr = inet_addr(LOCAL_DNS_IP);
    localAddr.sin_port = htons(DNS_PORT);
    if (bind(sockLocal, (SOCKADDR *) &localAddr, sizeof(SOCKADDR))) {
        printf("sockLocal binding failed!\n");
        exit(1);
    }

    sockServer = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.S_un.S_addr = inet_addr(SERVER_DNS_IP);
    serverAddr.sin_port = htons(DNS_PORT);

    printf("Socket initialize success!\n");

    //设置接收时限，每次套接字接收时限为timeout，如果超时则立即返回-1，重新进行下一次接收
    int timeout = 100;
    setsockopt(sockLocal, SOL_SOCKET, SO_RCVTIMEO, (const char *) &timeout, sizeof(timeout));
    setsockopt(sockServer, SOL_SOCKET, SO_RCVTIMEO, (const char *) &timeout, sizeof(timeout));

    while (1) {
        int len = sizeof(clientAddr);
        memset(recvBuf, 0, BUF_SIZE);
        //接受DNS请求, recvLen为接收到的字节数
        int recvLen = recvfrom(sockLocal, recvBuf, BUF_SIZE, 0, (SOCKADDR *) &clientAddr, &len);
        //接收失败
        if (recvLen == -1) {
//            printf("Receiving failed: %d\n", WSAGetLastError());
            continue;
        }
            //接收字符为0
        else if (recvLen == 0) {
            continue;
        }
            //正常接收字符
        else {
            getURL(recvBuf, recvLen);//获取域名
            int index = find(entries_num);//获取该域名在解析表中的位置

            //在域名解析表中没有找到
            if (index == -1) {
                //ID转换
                unsigned short *pID = (unsigned short *) malloc(sizeof(unsigned short));
                memcpy(pID, recvBuf, sizeof(unsigned short));//接收到的信息的前两个字节为ID信息
                unsigned short newID = htons(registerNewID(ntohs(*pID), clientAddr));
                memcpy(recvBuf, &newID, sizeof(newID));//将接收缓存中的老ID转换为新ID

                display(pID, newID, index, level, clientAddr, recvBuf);
                free(pID);

                //将recvBuf转发至外部DNS服务器
                if (sendto(sockServer, recvBuf, recvLen, 0, (SOCKADDR *) &serverAddr, sizeof(serverAddr)) == -1) {
                    printf("Fail to sendto sockServer: %d\n", WSAGetLastError());
                    continue;
                }

                //接收来自外部DNS服务器的响应报文
                recvLen = recvfrom(sockServer, recvBuf, BUF_SIZE, 0, (SOCKADDR *) &clientAddr, &len);
                if (recvLen == -1) {
                    printf("Fail to receive from sockServer: %d\n", WSAGetLastError());
                    continue;
                }

                //再次进行ID转换
                pID = (unsigned short *) malloc(sizeof(unsigned short));
                memcpy(pID, recvBuf, sizeof(unsigned short));//接收到的信息的前两个字节为ID信息
                int ID_index = ntohs(*pID);//接收到的ID即为新ID, 为ID转换表的下标
                unsigned short oldID = htons(IDTransTable[ID_index].oldID);
                memcpy(recvBuf, &oldID, sizeof(oldID));//将缓存中的ID换回老ID
                clientAddr = IDTransTable[ID_index].client;//更新获取客户端地址信息

                free(pID);

                //再把recvBuf转发回客户端处
                if (sendto(sockLocal, recvBuf, recvLen, 0, (SOCKADDR *) &clientAddr, sizeof(clientAddr)) ==
                    -1) {
                    printf("Fail to sendto client: %d\n", WSAGetLastError());
                    continue;
                }
            }
                //在域名解析表中找到了
            else {
                /*构造响应报文*/
                // 修改标志域
                memcpy(sendBuf, recvBuf, recvLen);
                unsigned short flag_field = htons(0x8180);//QR=1表示响应报文,RD=1,RA=1表示支持处理递归查询,其余标志位为0
                memcpy(&sendBuf[2], &flag_field, sizeof(flag_field));

                display(NULL, 0, index, level, clientAddr, recvBuf);

                // 修改回答数域
                unsigned short anCount_field;
                if (strcmp(DNS_table[index].ip, "0.0.0.0") == 0)
                    anCount_field = htons(0x0000);//屏蔽功能，回答数为0
                else
                    anCount_field = htons(0x0001);//服务器功能，回答数为1
                memcpy(&sendBuf[6], &anCount_field, sizeof(anCount_field));

                //构造DNS响应部分
                char answer[16];

                // 指针法标识域名,指向请求部分的域名
                // 域名固定出现的位置为包头开始偏移量为12字节的位置(前两位必定为11作为标识,后14位表示偏移量),故可用0xc00c指代域名位置
                unsigned short NAME = htons(0xc00c);
                memcpy(answer, &NAME, sizeof(NAME));

                unsigned short TYPE = htons(0x0001);//A地址类型
                memcpy(&answer[2], &TYPE, sizeof(TYPE));

                unsigned short CLASS = htons(0x0001);//IN(1), 指Internet数据
                memcpy(&answer[4], &CLASS, sizeof(CLASS));

                unsigned long TTL = htons(0x000000FF);//生存时间, 随便设置的
                memcpy(&answer[6], &TTL, sizeof(TTL));

                unsigned short RDLENGHT = htons(0x0004);//A类型地址是4个字节
                memcpy(&answer[10], &RDLENGHT, sizeof(RDLENGHT));

                unsigned long RDATA = (unsigned long) inet_addr(DNS_table[index].ip);//IP地址, 4个字节
                memcpy(&answer[12], &RDATA, sizeof(RDATA));

                //将修改后的请求报文和应答部分共同存入sendBuf
                memcpy(sendBuf + recvLen, answer, 16);

                //发送DNS响应报文
                if (sendto(sockLocal, sendBuf, sizeof(sendBuf), 0, (SOCKADDR *) &clientAddr, sizeof(clientAddr)) ==
                    -1) {
                    printf("Fail to sendto client: %d\n", WSAGetLastError());
                    continue;
                }
            }
        }
    }

    closesocket(sockLocal);
    closesocket(sockServer);
    WSACleanup();
    return 0;
}
